package com.tianji.learning.executor;

import com.tianji.common.utils.UserContext;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.stereotype.Component;

import java.util.concurrent.CompletableFuture;
import java.util.function.Supplier;
import java.util.concurrent.Executor;

/**
 * 线程池装饰器
 * 每个需要用户上下文的异步线程都应该手动设置 UserContext.setUser(currentUserId)
 *
 * @author Covey
 */
@Component
@Slf4j
public class ContextAwareExecutor {

    private final Executor asyncExecutor;

    public ContextAwareExecutor(@Qualifier("queryTaskExecutor") Executor asyncExecutor) {
        this.asyncExecutor = asyncExecutor;
    }

    public <T> CompletableFuture<T> supplyAsync(Supplier<T> supplier) {
        // 捕获当前（主线程）上下文
        Long userId = UserContext.getUser();

        return CompletableFuture.supplyAsync(() -> {
            try {
                // 恢复上下文
                if (userId != null) {
                    // 设置到新线程的ThreadLocal
                    UserContext.setUser(userId);
                }
                return supplier.get();
            } finally {
                // 清理上下文，避免内存泄漏
                UserContext.removeUser();
            }
        }, asyncExecutor);
    }
}